<?php

/**
 * This file is part of Webman AI.
 *
 * @author    walkor<walkor@workerman.net>
 * @copyright walkor<walkor@workerman.net>
 * @link      http://www.workerman.net/
 */

namespace plugin\ai\app\controller;

use plugin\ai\api\Chat;
use plugin\ai\app\event\data\ChatStreamData;
use plugin\ai\app\event\data\ModelRequestData;
use plugin\ai\app\event\data\ModelResponseData;
use plugin\ai\app\service\Common;
use plugin\ai\app\service\Setting;
use plugin\ai\app\service\User;
use support\exception\BusinessException;
use support\Request;
use support\Response;
use Webman\Event\Event;
use support\Log;
use Workerman\Protocols\Http\Chunk;

class ChatController extends Base
{

    /**
     * 不支持流式的模型
     *
     * @var string[]
     */
    protected $modelsNotSupportStream = ['o1'];

    /**
     * 不需要登录的方法
     *
     * @var string[]
     */
    protected $noNeedLogin = ['completions', 'translate', 'summarize'];

    /**
     * 对话
     * @param Request $request
     * @return Response
     * @throws BusinessException
     */
    public function completions(Request $request): Response
    {
        $model = $request->post('model');
        if ($error = static::tryReduceBalance($model)) {
            return json(['error' => ['message' => $error]]);
        }

        $stream = (bool)$request->post('stream');
        $messages = $request->post('messages');
        $content = last($messages)['content'] ?? '';
        $temperature = (float)$request->post('temperature', 0.5);
        $userMessageId = $request->post('user_message_id');
        $assistantMessageId = $request->post('assistant_message_id');
        $roleId = $request->post('role_id');
        $chatId = $request->post('chat_id');
        $userId = session('user.id') ?? session('user.uid');
        $sessionId = $request->sessionId();
        $realIp = $request->getRealIp();
        $this->saveMessage($userId, $sessionId, $userMessageId, $roleId, 'user', $content, $realIp, $model, $chatId);
        foreach ($this->modelsNotSupportStream as $notSupportModel) {
            if (stripos($model, $notSupportModel) === 0) {
                $stream = false;
                $temperature = 1;
                break;
            }
        }

        $connection = $request->connection;
        $modelRequestData = new ModelRequestData();
        $modelRequestData->data = [
            'model' => $model,
            'stream' => $stream,
            'messages' => $messages,
            'temperature' => $temperature,
        ];
        $modelRequestData->options = [
            'complete' => function ($result, $response) use ($connection, $userId, $sessionId, $assistantMessageId, $roleId, $realIp, $model, $chatId, $modelRequestData, $stream) {
                $responseData = new ModelResponseData();
                $responseData->data = $result;
                $responseData->modelRequestData = $modelRequestData;
                Event::dispatch('ai.chat.completions.response', $responseData);
                $result = $responseData->data;
                if (!$stream && is_string($result)) {
                    $connection->send(new Chunk(json_encode(['content' => $result], JSON_UNESCAPED_UNICODE) . "\n"));
                }
                if (isset($result['error'])) {
                    $result = json_encode($result, JSON_UNESCAPED_UNICODE);
                    Log::error($result);
                    $connection->send(new Chunk($result));
                    $modelType = Common::getModelType($model);
                    if ($userId && User::getBalance($userId, $modelType) > 0) {
                        User::addBalance($userId, $modelType);
                    }
                }
                $connection->send(new Chunk(''));
                $model = $responseData->modelRequestData->data['model'] ?? $model;
                $this->saveMessage($userId, $sessionId, $assistantMessageId, $roleId, 'assistant', $result, $realIp, $model, $chatId);
            }
        ];
        if ($stream) {
            $modelRequestData->options['stream'] = function ($data) use ($connection, $modelRequestData) {
                $chatStreamData = new ChatStreamData($modelRequestData, $data);
                Event::dispatch('ai.chat.completions.stream', $chatStreamData);
                $connection->send(new Chunk( json_encode($chatStreamData->streamData, JSON_UNESCAPED_UNICODE) . "\n"));
            };
        }
        Event::dispatch('ai.chat.completions.request', $modelRequestData);
        Chat::completions($modelRequestData->data, $modelRequestData->options);
        return response()->withHeaders([
            "Content-Type" => "application/json",
            "Transfer-Encoding" => "chunked",
        ]);
    }

    public function translate(Request $request)
    {
        $content = $request->post('content');
        if (!$content) {
            return json(['error' => ['message' => '内容为空']]);
        }
        $setting = Setting::getSetting();
        $model = $setting['translate_model'] ?? 'gpt-3.5-turbo';
        $prompt = !empty($setting['translate_prompt']) ? $setting['translate_prompt'] : '请作为翻译官，将以下内容翻译成英文。注意只做翻译，不回复其它内容。例如“请翻译:画一只小狗”你应该返回“a dog”；例如“请翻译:你能画一只熊猫么?”，你应该返回"a panda"；例如“请翻译:一只猫 --ar:1:2 --s 200”你应该返回“a cat --ar:1:2 --s 200”。记住你的返回中无论如何不能包含中文';
        $connection = $request->connection;
        $modelRequestData = new ModelRequestData();
        $modelRequestData->data = [
            'model' => $model,
            'temperature' => 1,
            'messages' => [
                ['role' => 'user', 'content' => $prompt],
                ['role' => 'assistant', 'content' => '好的，请发送要翻译的内容'],
                ['role' => 'user', 'content' => "请翻译:$content"]
            ],
        ];
        $modelRequestData->options = [
            'complete' => function ($result) use ($connection, $modelRequestData) {
                if (isset($result['error'])) {
                    $result = json_encode($result, JSON_UNESCAPED_UNICODE);
                    Log::error($result);
                    $connection->send(new Chunk($result));
                } else {
                    $responseData = new ModelResponseData();
                    $responseData->data = $result;
                    $responseData->modelRequestData = $modelRequestData;
                    Event::dispatch('ai.chat.translate.response', $responseData);
                    $connection->send(new Chunk(json_encode(['content' => $result])));
                }
                $connection->send(new Chunk(''));
            }
        ];
        Event::dispatch('ai.chat.translate.request', $modelRequestData);
        Chat::completions($modelRequestData->data, $modelRequestData->options);
        return response()->withHeaders([
            "Content-Type" => "application/json",
            "Transfer-Encoding" => "chunked",
        ]);
    }

    /**
     * 总结问题
     * @param Request $request
     * @return Response|null
     */
    public function summarize(Request $request): ?Response
    {
        $messages = $request->post('messages');
        $content = last($messages)['content'] ?? '';
        if (!$content) {
            return json(['error' => ['message' => '内容为空']]);
        }
        $connection = $request->connection;
        $model = $request->post('model');
        $content = json_encode($messages, JSON_UNESCAPED_UNICODE);
        $lastContent = last($messages)['content'];
        $content = $content . "\n\n" . '以上是一段对话，请将对话中最后一句总结成一个具有清晰含义的提问，注意，是总结成一个有明确含义的提问。如果无法总结，请返回 ' . $lastContent;
        Chat::completions([
            'model' => $model,
            'temperature' => 1,
            'messages' => [
                ['role' => 'user', 'content' => $content],
            ]], [
            'complete' => function ($result, $response) use ($connection) {
                if (isset($result['error'])) {
                    $result = json_encode($result, JSON_UNESCAPED_UNICODE);
                    Log::error($result);
                    $connection->send(new Chunk($result));
                } else {
                    $connection->send(new Chunk(json_encode(['content' => $result])));
                }
                $connection->send(new Chunk(''));
            }
        ]);
        return response()->withHeaders([
            "Content-Type" => "application/json",
            "Transfer-Encoding" => "chunked",
        ]);
    }

}
